import type { SimplifiedSchema } from 'mongodb-schema';
import {
  type PreferencesAccess,
  isAIFeatureEnabled,
} from 'compass-preferences-model/provider';
import type { AtlasService } from '@mongodb-js/atlas-service/provider';
import { AtlasServiceError } from '@mongodb-js/atlas-service/renderer';
import type { ConnectionInfo } from '@mongodb-js/connection-info';
import type { Document } from 'mongodb';
import type { Logger } from '@mongodb-js/compass-logging';
import { EJSON } from 'bson';
import { z } from 'zod';
import { getStore } from './store/atlas-ai-store';
import { optIntoGenAIWithModalPrompt } from './store/atlas-optin-reducer';
import {
  AtlasAiServiceInvalidInputError,
  AtlasAiServiceApiResponseParseError,
} from './atlas-ai-errors';

type GenerativeAiInput = {
  userInput: string;
  collectionName: string;
  databaseName: string;
  schema?: SimplifiedSchema;
  sampleDocuments?: Document[];
  signal: AbortSignal;
  requestId: string;
};

// The size/token validation happens on the server, however, we do
// want to ensure we're not uploading massive documents (some folks have documents > 1mb).
const AI_MAX_REQUEST_SIZE = 5120000;
const AI_MIN_SAMPLE_DOCUMENTS = 1;

type AIAggregation = {
  content: {
    aggregation?: {
      pipeline?: string;
    };
  };
};

type AIFeatureEnablement = {
  features: {
    [featureName: string]: {
      enabled: boolean;
    };
  };
};

type AIQuery = {
  content: {
    query: Record<
      'filter' | 'project' | 'collation' | 'sort' | 'skip' | 'limit',
      string
    >;
    aggregation?: { pipeline: string };
  };
};

function buildQueryOrAggregationMessageBody(
  input: Omit<GenerativeAiInput, 'signal' | 'requestId'>
) {
  const sampleDocuments = input.sampleDocuments
    ? EJSON.serialize(input.sampleDocuments, {
        relaxed: false,
      })
    : undefined;

  let msgBody = JSON.stringify({
    ...input,
    sampleDocuments,
  });
  if (msgBody.length > AI_MAX_REQUEST_SIZE && sampleDocuments) {
    // When the message body is over the max size, we try
    // to see if with fewer sample documents we can still perform the request.
    // If that fails we throw an error indicating this collection's
    // documents are too large to send to the ai.
    msgBody = JSON.stringify({
      ...input,
      sampleDocuments: EJSON.serialize(
        input.sampleDocuments?.slice(0, AI_MIN_SAMPLE_DOCUMENTS) || [],
        {
          relaxed: false,
        }
      ),
    });
  }

  if (msgBody.length > AI_MAX_REQUEST_SIZE) {
    throw new Error(
      'Sorry, your request is too large. Please use a smaller prompt or try using this feature on a collection with smaller documents.'
    );
  }

  return msgBody;
}

function hasExtraneousKeys(obj: any, expectedKeys: string[]) {
  return Object.keys(obj).some((key) => !expectedKeys.includes(key));
}

export function validateAIQueryResponse(
  response: any
): asserts response is AIQuery {
  const { content } = response ?? {};

  if (typeof content !== 'object' || content === null) {
    throw new Error('Unexpected response: expected content to be an object');
  }

  if (hasExtraneousKeys(content, ['query', 'aggregation'])) {
    throw new Error(
      'Unexpected keys in response: expected query and aggregation'
    );
  }

  const { query, aggregation } = content;

  if (!query && !aggregation) {
    throw new Error(
      'Unexpected response: expected query or aggregation, got none'
    );
  }

  if (query && typeof query !== 'object') {
    throw new Error('Unexpected response: expected query to be an object');
  }

  if (
    hasExtraneousKeys(query, [
      'filter',
      'project',
      'collation',
      'sort',
      'skip',
      'limit',
    ])
  ) {
    throw new Error(
      'Unexpected keys in response: expected filter, project, collation, sort, skip, limit, aggregation'
    );
  }

  for (const field of [
    'filter',
    'project',
    'collation',
    'sort',
    'skip',
    'limit',
  ]) {
    if (query[field] && typeof query[field] !== 'string') {
      throw new Error(
        `Unexpected response: expected field ${field} to be a string, got ${JSON.stringify(
          query[field],
          null,
          2
        )}`
      );
    }
  }

  if (aggregation && typeof aggregation.pipeline !== 'string') {
    throw new Error(
      `Unexpected response: expected aggregation pipeline to be a string, got ${JSON.stringify(
        aggregation,
        null,
        2
      )}`
    );
  }
}

export function validateAIAggregationResponse(
  response: any
): asserts response is AIAggregation {
  const { content } = response;

  if (typeof content !== 'object' || content === null) {
    throw new Error('Unexpected response: expected content to be an object');
  }

  if (hasExtraneousKeys(content, ['aggregation'])) {
    throw new Error('Unexpected keys in response: expected aggregation');
  }

  if (content.aggregation && typeof content.aggregation.pipeline !== 'string') {
    // Compared to queries where we will always get the `query` field, for
    // aggregations backend deletes the whole `aggregation` key if pipeline is
    // empty, so we only validate `pipeline` key if `aggregation` key is present
    throw new Error(
      `Unexpected response: expected aggregation to be a string, got ${String(
        content.aggregation.pipeline
      )}`
    );
  }
}

const aiURLConfig = {
  // There are two different sets of endpoints we use for our requests.
  // Down the line we'd like to only use the admin api, however,
  // we cannot currently call that from the Atlas UI. Pending CLOUDP-251201
  // NOTE: The unauthenticated endpoints are also rate limited by IP address
  // rather than by logged in user.
  'admin-api': {
    aggregation: 'unauth/ai/api/v1/mql-aggregation',
    query: 'unauth/ai/api/v1/mql-query',
  },
  cloud: {
    aggregation: (projectId: string) => {
      return `ai/v1/groups/${projectId}/mql-aggregation`;
    },
    query: (projectId: string) => {
      return `ai/v1/groups/${projectId}/mql-query`;
    },
    'mock-data-schema': (projectId: string) => {
      return `ai/v1/groups/${projectId}/mock-data-schema`;
    },
  },
} as const;

export interface MockDataSchemaRawField {
  type: string;
  sampleValues?: unknown[];
}

export interface MockDataSchemaRequest {
  collectionName: string;
  databaseName: string;
  schema: Record<string, MockDataSchemaRawField>;
  validationRules?: Record<string, unknown> | null;
  includeSampleValues?: boolean;
  requestId: string;
  signal: AbortSignal;
}

export const MockDataSchemaResponseShape = z.object({
  fields: z.array(
    z.object({
      fieldPath: z.string(),
      fakerMethod: z.string(),
      fakerArgs: z.array(
        z.union([
          z.object({
            json: z.string(),
          }),
          z.string(),
          z.number(),
          z.boolean(),
        ])
      ),
    })
  ),
});

export type MockDataSchemaResponse = z.infer<
  typeof MockDataSchemaResponseShape
>;

/**
 * The type of resource from the natural language query REST API
 */
type AIResourceType = 'query' | 'aggregation' | 'mock-data-schema';

export class AtlasAiService {
  private initPromise: Promise<void> | null = null;

  private apiURLPreset: 'admin-api' | 'cloud';
  private atlasService: AtlasService;
  private preferences: PreferencesAccess;
  private logger: Logger;

  constructor({
    apiURLPreset,
    atlasService,
    preferences,
    logger,
  }: {
    apiURLPreset: 'admin-api' | 'cloud';
    atlasService: AtlasService;
    preferences: PreferencesAccess;
    logger: Logger;
  }) {
    this.apiURLPreset = apiURLPreset;
    this.atlasService = atlasService;
    this.preferences = preferences;
    this.logger = logger;

    this.initPromise = this.setupAIAccess();
  }

  /**
   * @throws {AtlasAiServiceInvalidInputError} when given invalid arguments
   */
  private getUrlForEndpoint(
    resourceType: AIResourceType,
    { atlasMetadata }: ConnectionInfo
  ) {
    if (atlasMetadata) {
      return this.atlasService.cloudEndpoint(
        aiURLConfig.cloud[resourceType](atlasMetadata.projectId)
      );
    }

    if (resourceType === 'mock-data-schema') {
      throw new AtlasAiServiceInvalidInputError(
        "Can't perform generative ai request: mock-data-schema is not available for admin-api"
      );
    }

    const urlPath = aiURLConfig['admin-api'][resourceType];
    return this.atlasService.adminApiEndpoint(urlPath);
  }

  private throwIfAINotEnabled() {
    if (process.env.COMPASS_E2E_SKIP_AI_OPT_IN === 'true') {
      return;
    }
    if (!isAIFeatureEnabled(this.preferences.getPreferences())) {
      throw new Error(
        "Compass' AI functionality is not currently enabled. Please try again later."
      );
    }
  }

  async setupAIAccess(): Promise<void> {
    // We default GEN_AI_ACCESS on for everyone. Down the line if/when
    // we add more features with partial rollout, we'll fetch access here.
    await this.preferences.savePreferences({
      cloudFeatureRolloutAccess: {
        GEN_AI_COMPASS: true,
      },
    });
  }

  async ensureAiFeatureAccess({ signal }: { signal?: AbortSignal } = {}) {
    return getStore().dispatch(
      optIntoGenAIWithModalPrompt({
        signal,
        isCloudOptIn: this.apiURLPreset === 'cloud',
      })
    );
  }

  private getQueryOrAggregationFromUserInput = async <T>(
    {
      urlId,
      input,
      connectionInfo,
    }: {
      urlId: 'query' | 'aggregation';
      input: GenerativeAiInput;
      connectionInfo: ConnectionInfo;
    },
    validationFn: (res: any) => asserts res is T
  ): Promise<T> => {
    await this.initPromise;
    this.throwIfAINotEnabled();

    const { signal, requestId, ...rest } = input;
    const msgBody = buildQueryOrAggregationMessageBody(rest);

    const url = `${this.getUrlForEndpoint(
      urlId,
      connectionInfo
    )}?request_id=${encodeURIComponent(requestId)}`;

    this.logger.log.info(
      this.logger.mongoLogId(1_001_000_308),
      'AtlasAIService',
      'Running AI query generation request',
      {
        url,
        userInput: input.userInput,
        collectionName: input.collectionName,
        databaseName: input.databaseName,
        messageBodyLength: msgBody.length,
        requestId,
      }
    );

    const res = await this.atlasService.authenticatedFetch(url, {
      signal,
      method: 'POST',
      body: msgBody,
      headers: {
        'Content-Type': 'application/json',
        Accept: 'application/json',
      },
    });

    // Sometimes the server will return empty response and calling res.json directly
    // throws and user see "Unexpected end of JSON input" error, which is not helpful.
    // So we will get the text from the response first and then try to parse it.
    // If it fails, we will throw a more helpful error message.
    const text = await res.text();
    let data;
    try {
      data = JSON.parse(text);
    } catch {
      this.logger.log.info(
        this.logger.mongoLogId(1_001_000_310),
        'AtlasAIService',
        'Failed to parse the response from AI API',
        {
          text,
          requestId,
        }
      );
      throw new AtlasServiceError(
        'ServerError',
        500, // Not using res.status as its 200 in this case
        'Internal server error',
        'INTERNAL_SERVER_ERROR'
      );
    }
    validationFn(data);
    return data;
  };

  async getAggregationFromUserInput(
    input: GenerativeAiInput,
    connectionInfo: ConnectionInfo
  ) {
    return this.getQueryOrAggregationFromUserInput(
      {
        connectionInfo,
        urlId: 'aggregation',
        input,
      },
      validateAIAggregationResponse
    );
  }

  async getQueryFromUserInput(
    input: GenerativeAiInput,
    connectionInfo: ConnectionInfo
  ) {
    return this.getQueryOrAggregationFromUserInput(
      {
        urlId: 'query',
        input,
        connectionInfo,
      },
      validateAIQueryResponse
    );
  }

  async getMockDataSchema(
    input: MockDataSchemaRequest,
    connectionInfo: ConnectionInfo
  ): Promise<MockDataSchemaResponse> {
    const { collectionName, databaseName } = input;
    let schema = input.schema;

    const url = `${this.getUrlForEndpoint(
      'mock-data-schema',
      connectionInfo
    )}?request_id=${encodeURIComponent(input.requestId)}`;

    if (!input.includeSampleValues) {
      const newSchema: Record<
        string,
        Omit<MockDataSchemaRawField, 'sampleValues'>
      > = {};
      for (const [k, v] of Object.entries(schema)) {
        newSchema[k] = { type: v.type };
      }
      schema = newSchema;
    }

    const res = await this.atlasService.authenticatedFetch(url, {
      method: 'POST',
      body: JSON.stringify({
        collectionName,
        databaseName,
        schema,
        validationRules: input.validationRules,
      }),
      headers: {
        'Content-Type': 'application/json',
        Accept: 'application/json',
      },
      signal: input.signal,
    });

    try {
      const data = await res.json();
      return MockDataSchemaResponseShape.parse(data);
    } catch (err) {
      const errorMessage = err instanceof Error ? err.stack : String(err);
      this.logger.log.error(
        this.logger.mongoLogId(1_001_000_311),
        'AtlasAiService',
        'Failed to parse mock data schema response with expected schema',
        {
          namespace: `${databaseName}.${collectionName}`,
          message: errorMessage,
        }
      );
      throw new AtlasAiServiceApiResponseParseError(
        'Response does not match expected schema'
      );
    }
  }

  async optIntoGenAIFeatures() {
    if (this.apiURLPreset === 'cloud') {
      // Performs a post request to Atlas to set the user opt in preference to true.
      await this.atlasService.authenticatedFetch(
        this.atlasService.cloudEndpoint(
          'settings/optInDataExplorerGenAIFeatures'
        ),
        {
          method: 'POST',
          headers: {
            'Content-Type': 'application/x-www-form-urlencoded',
            Accept: 'application/json',
          },
          body: new URLSearchParams([['value', 'true']]),
        }
      );
    }
    await this.preferences.savePreferences({
      optInGenAIFeatures: true,
    });
  }

  private validateAIFeatureEnablementResponse(
    response: any
  ): asserts response is AIFeatureEnablement {
    const { features } = response;
    if (typeof features !== 'object') {
      throw new Error('Unexpected response: expected features to be an object');
    }
  }
}
